package cn.genmer.test.security.machinelearning.deeplearning4j.mnist.V2.strategy.concrete;

import cn.genmer.test.security.machinelearning.deeplearning4j.mnist.V1.MnistTrain;
import cn.genmer.test.security.machinelearning.deeplearning4j.mnist.V2.strategy.MnistModelStrategy;
import org.deeplearning4j.nn.multilayer.MultiLayerNetwork;
import org.deeplearning4j.util.ModelSerializer;
import org.nd4j.linalg.api.ndarray.INDArray;
import org.nd4j.linalg.factory.Nd4j;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.imageio.ImageIO;
import java.awt.image.BufferedImage;
import java.io.File;
import java.io.IOException;

public class CnnStrategy implements MnistModelStrategy {
  private Logger logger = LoggerFactory.getLogger(CnnStrategy.class);

  private MultiLayerNetwork model;
  public CnnStrategy() {
    // 加载预训练好的LeNet模型
    File locationToSave = new File(MnistTrain.BASE_PATH +"/cnn_mnist_model.zip");
    System.out.println("预训练模型的路径：" + locationToSave.getAbsolutePath());
    try {
      model = ModelSerializer.restoreMultiLayerNetwork(locationToSave);
    } catch (IOException e) {
      e.printStackTrace();
    }
  }

  @Override
  public int predict(INDArray feature) {
    System.out.println("当前使用策略为：CNN: Convolutional Neural Network （卷积神经网络）");
    // 对图像进行推理，并返回预测结果的索引
    INDArray output = model.output(feature);
    return Nd4j.argMax(output, 1).getInt(0);
  }

  @Override
  public INDArray loadImgAndGetFeature(BufferedImage image) {

    int width = image.getWidth();
    int height = image.getHeight();

    // 将RGB图片变为灰度图，并将每个像素存于二维数组
    int[][] grayArray = new int[height][width];
    for (int i = 0; i < height; i++) {
      for (int j = 0; j < width; j++) {
        int rgb = image.getRGB(j, i);
        int gray = (int) (0.2989 * ((rgb >> 16) & 0xff) + 0.5870 * ((rgb >> 8) & 0xff) + 0.1140 * (rgb & 0xff));
        // 二值化
        if (gray <= 120)
          grayArray[i][j] = 255;
      }
    }

    // 对当前的二维数组resize，因为训练好的模型要求输入图像大小为28×28
    int[][] expandedSubArray = resizeImage2D(grayArray, 28, 28);

    // 创建INDArray对象
    INDArray array = getFeatures(expandedSubArray);
    // 调整形状并进行数值计算
    INDArray input = array.reshape(1, 1, 28, 28).divi(255.0);

    return input;
  }

  @Override
  public INDArray loadImgAndGetFeature(String modelPath) {
    try {
      BufferedImage img = ImageIO.read(new File(modelPath));
      return loadImgAndGetFeature(img);
    } catch (Exception e){
      logger.error("【图像加载异常】", e);
    }
    return null;
  }


  static INDArray getFeatures(int[][] inputImage){
    // 获取输入参数的形状
    int[] shape = getShape(inputImage);
    // 创建INDArray对象
    INDArray array = Nd4j.create(shape);

    // 将多维数组的值复制到INDArray对象中
    for (int i = 0; i < shape[0]; i++) {
      for (int j = 0; j < shape[1]; j++) {
        array.putScalar(new int[]{i, j}, inputImage[i][j]);
      }
    }
    return array;
  }

  static int[] getShape(int[][] array) {
    int firstDimensionSize = array.length;
    int secondDimensionSize = array[0].length;

    return new int[]{firstDimensionSize, secondDimensionSize};
  }


  private static void save(int[][] data, String path) {
    int width = data[0].length;
    int height = data.length;
    BufferedImage image = new BufferedImage(width, height, BufferedImage.TYPE_BYTE_GRAY);

    for (int y = 0; y < height; y++) {
      for (int x = 0; x < width; x++) {
        int gray = data[y][x];
        int rgb = (gray << 16) | (gray << 8) | gray;
        image.setRGB(x, y, rgb);
      }
    }

    try {
      ImageIO.write(image, "png", new File(path));
    } catch (Exception e) {
      e.printStackTrace();
    }
  }

  private static int[][] resizeImage2D(int[][] input, int newWidth, int newHeight) {
    int height = input.length;
    int width = input[0].length;
    float widthRatio = (float) width / newWidth;
    float heightRatio = (float) height / newHeight;
    int[][] output = new int[newHeight][newWidth];

    for (int y = 0; y < newHeight; y++) {
      for (int x = 0; x < newWidth; x++) {
        float px = x * widthRatio;
        float py = y * heightRatio;
        int x1 = (int) Math.floor(px);
        int y1 = (int) Math.floor(py);
        int x2 = Math.min(x1 + 1, width - 1);
        int y2 = Math.min(y1 + 1, height - 1);

        float dx = px - x1;
        float dy = py - y1;

        float topLeft = input[y1][x1];
        float topRight = input[y1][x2];
        float bottomLeft = input[y2][x1];
        float bottomRight = input[y2][x2];

        float topAvg = topLeft + dx * (topRight - topLeft);
        float bottomAvg = bottomLeft + dx * (bottomRight - bottomLeft);

        output[y][x] = (int) (topAvg + dy * (bottomAvg - topAvg));
      }
    }

    return output;
  }
}
